#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmssl/sm4.h>
#include <gmssl/rand.h>

#define BLOCK_SIZE 16

// PKCS#7 填充
void pkcs7_pad(const uint8_t *in, size_t in_len, uint8_t **out, size_t *out_len) {
    size_t pad_len = BLOCK_SIZE - (in_len % BLOCK_SIZE);
    *out_len = in_len + pad_len;
    *out = malloc(*out_len);
    memcpy(*out, in, in_len);
    for (size_t i = 0; i < pad_len; i++) {
        (*out)[in_len + i] = (uint8_t)pad_len;
    }
}

// PKCS#7 去填充
int pkcs7_unpad(uint8_t *data, size_t data_len, size_t *unpadded_len) {
    if (data_len == 0 || data_len % BLOCK_SIZE != 0) return 0;
    uint8_t pad = data[data_len - 1];
    if (pad < 1 || pad > BLOCK_SIZE) return 0;
    for (size_t i = 0; i < pad; i++) {
        if (data[data_len - 1 - i] != pad) return 0;
    }
    *unpadded_len = data_len - pad;
    return 1;
}

// 手动实现 SM4-CBC 加密
void sm4_cbc_encrypt_manual(const uint8_t *in, uint8_t *out, size_t len, const SM4_KEY *key, uint8_t iv[16]) {
    uint8_t prev[16];
    memcpy(prev, iv, 16);

    for (size_t i = 0; i < len; i += BLOCK_SIZE) {
        // XOR with previous block (or IV for first block)
        uint8_t input[16];
        for (int j = 0; j < BLOCK_SIZE; j++) {
            input[j] = in[i + j] ^ prev[j];
        }
        // Encrypt
        sm4_encrypt(key, input, &out[i]);
        // Update prev
        memcpy(prev, &out[i], 16);
    }
}

// 手动实现 SM4-CBC 解密
void sm4_cbc_decrypt_manual(const uint8_t *in, uint8_t *out, size_t len, const SM4_KEY *key, uint8_t iv[16]) {
    uint8_t prev[16];
    memcpy(prev, iv, 16);

    for (size_t i = 0; i < len; i += BLOCK_SIZE) {
        // Save current ciphertext block (needed for next XOR)
        uint8_t curr[16];
        memcpy(curr, &in[i], 16);
        // Decrypt
        sm4_encrypt(key, &in[i], &out[i]); 
        // XOR with previous ciphertext (or IV)
        for (int j = 0; j < BLOCK_SIZE; j++) {
            out[i + j] ^= prev[j];
        }
        // Update prev
        memcpy(prev, curr, 16);
    }
}

int main() {
    // 读取 sn.txt
    FILE *fp = fopen("sn.txt", "rb");
    if (!fp) {
        perror("fopen sn.txt");
        return 1;
    }
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    uint8_t *plaintext = malloc(file_size);
    fread(plaintext, 1, file_size, fp);
    fclose(fp);

    // 生成随机密钥和 IV（使用 GmSSL 的 rand_bytes）
    uint8_t key[16], iv[16];
    if (rand_bytes(key, 16) != 1 || rand_bytes(iv, 16) != 1) {
        fprintf(stderr, "Failed to generate random key/IV\n");
        free(plaintext);
        return 1;
    }

    // 填充明文
    uint8_t *padded_plaintext;
    size_t padded_len;
    pkcs7_pad(plaintext, file_size, &padded_plaintext, &padded_len);
    free(plaintext);

    // 分配缓冲区
    uint8_t *ciphertext = malloc(padded_len);
    uint8_t *decrypted = malloc(padded_len);

    // 设置密钥
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key);

    // CBC 加密（手动实现）
    sm4_cbc_encrypt_manual(padded_plaintext, ciphertext, padded_len, &sm4_key, iv);

    // 保存 IV + 密文
    FILE *enc_fp = fopen("sn.enc", "wb");
    fwrite(iv, 1, 16, enc_fp);
    fwrite(ciphertext, 1, padded_len, enc_fp);
    fclose(enc_fp);

    // 解密：读取 IV 和密文
    FILE *dec_fp = fopen("sn.enc", "rb");
    uint8_t dec_iv[16];
    fread(dec_iv, 1, 16, dec_fp);
    size_t enc_data_len = padded_len;
    uint8_t *enc_data = malloc(enc_data_len);
    fread(enc_data, 1, enc_data_len, dec_fp);
    fclose(dec_fp);

    // 设置解密密钥
    sm4_set_decrypt_key(&sm4_key, key);

    // CBC 解密（手动实现）
    sm4_cbc_decrypt_manual(enc_data, decrypted, enc_data_len, &sm4_key, dec_iv);
    free(enc_data);

    // 去填充
    size_t unpadded_len;
    if (!pkcs7_unpad(decrypted, enc_data_len, &unpadded_len)) {
        fprintf(stderr, "Padding error\n");
        free(padded_plaintext);
        free(ciphertext);
        free(decrypted);
        return 1;
    }

    // 写入 sn2.txt
    FILE *out_fp = fopen("sn2.txt", "wb");
    fwrite(decrypted, 1, unpadded_len, out_fp);
    fclose(out_fp);

    printf("SM4-CBC encryption and decryption completed successfully.\n");
    printf("Original size: %ld bytes\n", file_size);
    printf("Decrypted size: %zu bytes\n", unpadded_len);

    // 清理内存
    free(padded_plaintext);
    free(ciphertext);
    free(decrypted);

    return 0;
}
